import onnx
import numpy as np

onnx_model = onnx.load('yolov5m.onnx')
graph = onnx_model.graph

for i in range(40,-1,-1):
    node = graph.node[i]
    graph.node.remove(node)

weight=np.array(
    [[[[1,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[1,0],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[0,0]],[[1,0],[0,0]]],

    [[[0,1],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,1],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[0,0]],[[0,1],[0,0]]],

    [[[0,0],[1,0]],[[0,0],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[1,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[1,0]]],

    [[[0,0],[0,1]],[[0,0],[0,0]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[0,1]],[[0,0],[0,0]]],
    [[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,1]]]],dtype=np.float32)

SliceConvWeight = onnx.helper.make_tensor('SliceConvWeight', onnx.TensorProto.FLOAT, [12,3,2,2], weight.flatten())
graph.initializer.append(SliceConvWeight)
new_node = onnx.helper.make_node(
    'Conv',
    name='SliceConv',
    inputs=['images', 'SliceConvWeight'],
    outputs=['SliceOut'],
    dilations=[1,1],
    group=1,
    kernel_shape=[2,2],
    pads=[0,0,0,0],
    strides=[2,2]
)
graph.node.insert(0,new_node)

for node in graph.node:
    if node.name=='Conv_41': # 原slice后接的conv的名字
        print(node)
        node.input[0]='SliceOut'


onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'yolov5m_modify.onnx')